import gurobipy as gp
from gurobipy import GRB
import numpy as np

from DSB.SWT1.Code.algorithm_comparison import GasSchedulingInstanceGenerator


class GasScheduleGurobi:
    def __init__(self, instance):
        # 加载数据（同之前的算例格式）
        self.params = instance['parameters']
        self.processing_times = instance['processing_times']

        # 基本参数
        self.num_wp = self.params['num_workpieces']  # 气瓶数量
        self.num_comp = self.params['num_components']  # 组分数量
        self.num_inflate_eq = self.params['num_inflation_eq']  # 充气设备数
        self.num_analysis_eq = self.params['num_analysis_eq_per_component']  # 每组分分析设备数
        self.num_clean_eq = self.params['num_clean_eq']  # 清洗设备数
        self.num_shake_eq = self.params['num_shake_eq']  # 摇匀设备数

        # 时间参数
        self.clean_time = self.processing_times['clean_time']  # 清洗时间（固定）
        self.inflate_time = np.array(self.processing_times['inflation_time'])  # 充气时间[气瓶][组分]
        self.shake_time = np.array([self.num_comp * self.processing_times['shake_unit_time']
                                    for _ in range(self.num_wp)])  # 摇匀时间[气瓶]
        self.analysis_time = np.array(self.processing_times['analysis_time'])  # 分析时间[气瓶][组分]

        # 创建模型
        self.model = gp.Model("Gas_Scheduling")

        # 决策变量（核心）
        self._init_variables()

        # 添加约束
        self._add_constraints()

        # 目标函数：最小化总完成时间
        self.model.setObjective(self.C_max, GRB.MINIMIZE)

    def _init_variables(self):
        # 1. 清洗阶段
        self.clean_s = self.model.addVars(self.num_wp, vtype=GRB.CONTINUOUS, name="clean_start")  # 开始时间
        self.clean_e = self.model.addVars(self.num_wp, vtype=GRB.CONTINUOUS, name="clean_end")  # 结束时间
        self.clean_eq = self.model.addVars(self.num_wp, self.num_clean_eq, vtype=GRB.BINARY, name="clean_eq")  # 设备分配

        # 2. 充气阶段（按组分依次充气）
        self.inflate_s = self.model.addVars(self.num_wp, self.num_comp, vtype=GRB.CONTINUOUS,
                                            name="inflate_start")  # 开始时间
        self.inflate_e = self.model.addVars(self.num_wp, self.num_comp, vtype=GRB.CONTINUOUS,
                                            name="inflate_end")  # 结束时间
        self.inflate_eq = self.model.addVars(self.num_wp, self.num_comp, self.num_inflate_eq, vtype=GRB.BINARY,
                                             name="inflate_eq")  # 设备分配

        # 3. 摇匀阶段
        self.shake_s = self.model.addVars(self.num_wp, vtype=GRB.CONTINUOUS, name="shake_start")  # 开始时间
        self.shake_e = self.model.addVars(self.num_wp, vtype=GRB.CONTINUOUS, name="shake_end")  # 结束时间
        self.shake_eq = self.model.addVars(self.num_wp, self.num_shake_eq, vtype=GRB.BINARY, name="shake_eq")  # 设备分配

        # 4. 分析阶段（分组串联测试）
        self.group = self.model.addVars(self.num_wp, self.num_comp, self.num_comp, vtype=GRB.BINARY,
                                        name="group")  # 分组标记[气瓶][组分][组号]
        self.analysis_g_s = self.model.addVars(self.num_wp, self.num_comp, vtype=GRB.CONTINUOUS,
                                               name="group_start")  # 组开始时间
        self.analysis_g_e = self.model.addVars(self.num_wp, self.num_comp, vtype=GRB.CONTINUOUS,
                                               name="group_end")  # 组结束时间
        self.analysis_eq = self.model.addVars(self.num_wp, self.num_comp, self.num_analysis_eq, vtype=GRB.BINARY,
                                              name="analysis_eq")  # 设备分配

        # 总完成时间
        self.C_max = self.model.addVar(vtype=GRB.CONTINUOUS, name="C_max")

    def _add_constraints(self):
        # 1. 清洗阶段约束
        # 1.1 每个气瓶仅分配到一个清洗设备
        for wp in range(self.num_wp):
            self.model.addConstr(gp.quicksum(self.clean_eq[wp, eq] for eq in range(self.num_clean_eq)) == 1)
        # 1.2 清洗时间关系（结束=开始+时间）
        for wp in range(self.num_wp):
            self.model.addConstr(self.clean_e[wp] == self.clean_s[wp] + self.clean_time)
        # 1.3 清洗设备冲突约束（同一设备上的任务不重叠）
        for eq in range(self.num_clean_eq):
            for wp1 in range(self.num_wp):
                for wp2 in range(wp1 + 1, self.num_wp):
                    self.model.addConstr(
                        self.clean_s[wp1] >= self.clean_e[wp2] + 1e6 * (
                                    1 - self.clean_eq[wp1, eq] * self.clean_eq[wp2, eq])
                    )
                    self.model.addConstr(
                        self.clean_s[wp2] >= self.clean_e[wp1] + 1e6 * (
                                    1 - self.clean_eq[wp1, eq] * self.clean_eq[wp2, eq])
                    )

        # 2. 充气阶段约束
        # 2.1 每个气瓶的组分仅分配到一个充气设备
        for wp in range(self.num_wp):
            for comp in range(self.num_comp):
                self.model.addConstr(
                    gp.quicksum(self.inflate_eq[wp, comp, eq] for eq in range(self.num_inflate_eq)) == 1)
        # 2.2 充气时间关系
        for wp in range(self.num_wp):
            for comp in range(self.num_comp):
                self.model.addConstr(self.inflate_e[wp, comp] == self.inflate_s[wp, comp] + self.inflate_time[wp, comp])
        # 2.3 同一气瓶的组分需按顺序充气（后一组分需等前一组分完成）
        for wp in range(self.num_wp):
            for comp in range(1, self.num_comp):
                self.model.addConstr(self.inflate_s[wp, comp] >= self.inflate_e[wp, comp - 1])
        # 2.4 充气需在清洗完成后开始
        for wp in range(self.num_wp):
            for comp in range(self.num_comp):
                self.model.addConstr(self.inflate_s[wp, comp] >= self.clean_e[wp])
        # 2.5 充气设备冲突约束
        for eq in range(self.num_inflate_eq):
            for wp1 in range(self.num_wp):
                for comp1 in range(self.num_comp):
                    for wp2 in range(self.num_wp):
                        for comp2 in range(self.num_comp):
                            if (wp1 == wp2 and comp1 == comp2):
                                continue
                            self.model.addConstr(
                                self.inflate_s[wp1, comp1] >= self.inflate_e[wp2, comp2] + 1e6 * (
                                        1 - self.inflate_eq[wp1, comp1, eq] * self.inflate_eq[wp2, comp2, eq]
                                )
                            )

        # 3. 摇匀阶段约束
        # 3.1 每个气瓶仅分配到一个摇匀设备
        for wp in range(self.num_wp):
            self.model.addConstr(gp.quicksum(self.shake_eq[wp, eq] for eq in range(self.num_shake_eq)) == 1)
        # 3.2 摇匀时间关系
        for wp in range(self.num_wp):
            self.model.addConstr(self.shake_e[wp] == self.shake_s[wp] + self.shake_time[wp])
        # 3.3 摇匀需在充气完成后开始（等待最后一个组分充气完成）
        for wp in range(self.num_wp):
            self.model.addConstr(self.shake_s[wp] >= gp.quicksum(
                self.inflate_e[wp, comp] for comp in range(self.num_comp)
            ) / self.num_comp  # 取最大值的简化约束，实际需用max函数（可通过Big-M实现）
                                 )
        # 3.4 摇匀设备冲突约束（类似清洗设备）
        for eq in range(self.num_shake_eq):
            for wp1 in range(self.num_wp):
                for wp2 in range(wp1 + 1, self.num_wp):
                    self.model.addConstr(
                        self.shake_s[wp1] >= self.shake_e[wp2] + 1e6 * (
                                    1 - self.shake_eq[wp1, eq] * self.shake_eq[wp2, eq])
                    )

        # 4. 分析阶段约束（核心：串联组约束）
        # 4.1 每个组分属于且仅属于一个组
        for wp in range(self.num_wp):
            for comp in range(self.num_comp):
                self.model.addConstr(gp.quicksum(self.group[wp, comp, g] for g in range(self.num_comp)) == 1)
        # 4.2 组内所有组分同时开始、同时结束（结束时间=开始+组内最长分析时间）
        for wp in range(self.num_wp):
            for g in range(self.num_comp):
                # 组内是否有组分
                has_component = self.model.addVar(vtype=GRB.BINARY, name=f"has_comp_{wp}_{g}")
                # 组内是否有组分（求和≥1时，has_component=1）
                sum_group = gp.quicksum(self.group[wp, comp, g] for comp in range(self.num_comp))
                # 1. 如果sum_group ≥ 1，则has_component必须为1（sum_group ≥ 1 → has_component=1）
                self.model.addConstr(sum_group <= self.num_comp * has_component)  # 当has_component=0时，sum_group必须=0
                # 2. 如果has_component=1，则sum_group至少为1（has_component=1 → sum_group ≥ 1）
                self.model.addConstr(sum_group >= has_component)

                # 组开始时间=组内所有组分的开始时间
                for comp in range(self.num_comp):
                    self.model.addConstr(
                        self.analysis_g_s[wp, g] == self.analysis_g_s[wp, g] * (1 - self.group[wp, comp, g])
                        + self.analysis_g_s[wp, g] * self.group[wp, comp, g]
                    )

                # 组结束时间=开始时间+组内最长分析时间
                max_analysis = self.model.addVar(vtype=GRB.CONTINUOUS, name=f"max_analysis_{wp}_{g}")
                for comp in range(self.num_comp):
                    self.model.addConstr(
                        max_analysis >= self.analysis_time[wp, comp] * self.group[wp, comp, g]
                    )
                self.model.addConstr(self.analysis_g_e[wp, g] == self.analysis_g_s[wp, g] + max_analysis)

        # 4.3 分析需在摇匀完成后开始
        for wp in range(self.num_wp):
            for g in range(self.num_comp):
                self.model.addConstr(self.analysis_g_s[wp, g] >= self.shake_e[wp])

        # 5. 总完成时间约束（所有阶段最后结束时间）
        self.model.addConstr(self.C_max >= gp.quicksum(self.clean_e[wp] for wp in range(self.num_wp)) / self.num_wp)
        self.model.addConstr(self.C_max >= gp.quicksum(
            self.inflate_e[wp, comp] for wp in range(self.num_wp) for comp in range(self.num_comp))
                             / (self.num_wp * self.num_comp))
        self.model.addConstr(self.C_max >= gp.quicksum(self.shake_e[wp] for wp in range(self.num_wp)) / self.num_wp)
        self.model.addConstr(self.C_max >= gp.quicksum(
            self.analysis_g_e[wp, g] for wp in range(self.num_wp) for g in range(self.num_comp))
                             / (self.num_wp * self.num_comp))

    def optimize(self, time_limit=300000):
        self.model.setParam("TimeLimit", time_limit)  # 求解时间限制（秒）
        self.model.optimize()

        if self.model.status == GRB.OPTIMAL:
            print(f"最优总时间: {self.C_max.X:.2f}")
            # 可在此处提取各阶段的时间和设备分配结果
        else:
            print("未找到最优解")


# 示例：加载算例并求解
if __name__ == "__main__":
    # 假设instance是从JSON加载的算例数据（同之前的格式）
    instance = GasSchedulingInstanceGenerator.load_instance("gas_scheduling_instances/small/instance_1.json")
    model = GasScheduleGurobi(instance)
    model.optimize()